package com.kly.ai.rags;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * @author admin
 * @date 2025/5/29
 */
@Component
@Slf4j
@RequiredArgsConstructor
public class LoadDocumentLocal {

    private final ResourcePatternResolver resourcePatternResolver;

    /**
     * 加载love documents
     * @return
     */
    public List<Document> loadLoveDocuments() {
        List<Document> documents = new ArrayList<>();
        try {
            Resource[] resources = resourcePatternResolver.getResources("classpath:documents/*.md");
            getMD(documents, resources);
        } catch (IOException e) {
            log.error(e.getMessage());
        }

        return documents;
    }


    /**
     * 加载psychic documents
     */
    public List<Document> loadPsychicDocuments() {
        List<Document> documents = new ArrayList<>();
        try {
            Resource[] resources = resourcePatternResolver.getResources("classpath:documents/psychologist/*.md");
            getMD(documents, resources);
        } catch (IOException e) {
            log.error(e.getMessage());
            e.printStackTrace();
        }
        return documents;
    }

    /**
     * 转换文档
     * @param documents
     * @param resources
     */
    private void getMD(List<Document> documents, Resource[] resources) {
        for (Resource resource : resources) {
            String filename = resource.getFilename();
            MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                    .withHorizontalRuleCreateDocument(true)
                    .withIncludeCodeBlock(false)
                    .withIncludeBlockquote(false)
                    .withAdditionalMetadata("filename", filename)
                    .build();

            MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, config);
            documents.addAll(reader.read());
        }
    }
}
